import time
import os
import matplotlib.pyplot as plt
import numpy as np

# 导入你的加密算法实现
from publickey.rsa_core import RSA
from publickey.elgamal_core import ElGamal
from publickey.ecc_core import get_curve_by_name, generate_keypair as ecc_gen_keys, ecc_encrypt_simplified, ecc_decrypt_simplified

def run_single_performance_test(algo_type, config, num_iterations=100, message_size_bytes=40):
    """
    对单个算法配置进行性能测试的辅助函数。
    修改: 所有时间结果直接以毫秒(ms)为单位返回。
    """
    results = {}
    print(f"\n--- 测试 {algo_type} with config: {config} ---")
    messages = [os.urandom(message_size_bytes) for _ in range(num_iterations)]

    # --- 密钥生成 ---
    start_time = time.perf_counter()
    if algo_type == "RSA":
        instance = RSA(bits=config)
        priv_key, pub_key = instance.private_key, instance.public_key
    elif algo_type == "ECC":
        instance = get_curve_by_name(config)
        priv_key, pub_key = ecc_gen_keys(instance)
    elif algo_type == "ElGamal":
        instance = ElGamal(key_size_bits=config)
        priv_key, pub_key = instance.private_key, instance.public_key
    else:
        raise ValueError("不支持的算法类型")
    # 修改: 时间单位转换为毫秒
    results['key_gen_time'] = (time.perf_counter() - start_time) * 1000
    print(f"密钥生成时间: {results['key_gen_time']:.4f} 毫秒")
    
    # --- 加密 ---
    start_time = time.perf_counter()
    if algo_type == "RSA":
        ciphertexts = [instance.encrypt(m) for m in messages]
    elif algo_type == "ECC":
        ciphertexts = [ecc_encrypt_simplified(pub_key, instance, m) for m in messages]
    elif algo_type == "ElGamal":
        messages_int = [ElGamal.string_to_int(m.decode('latin-1')) for m in messages]
        ciphertexts = [ElGamal.encrypt(pub_key, m) for m in messages_int]
    results['encrypt_time'] = ((time.perf_counter() - start_time) / num_iterations) * 1000
    print(f"加密平均时间: {results['encrypt_time']:.8f} 毫秒")

    # --- 解密 ---
    start_time = time.perf_counter()
    if algo_type == "RSA":
        for c in ciphertexts: instance.decrypt(c)
    elif algo_type == "ECC":
        for R, C in ciphertexts: ecc_decrypt_simplified(priv_key, instance, R, C)
    elif algo_type == "ElGamal":
        public_params = (pub_key[0], pub_key[1], pub_key[2])
        for c in ciphertexts: ElGamal.decrypt(priv_key, public_params, c)
    results['decrypt_time'] = ((time.perf_counter() - start_time) / num_iterations) * 1000
    print(f"解密平均时间: {results['decrypt_time']:.8f} 毫秒")

    return results

def compare_all_at_128bit_security():
    """
    主函数：在128-bit等效安全级别下，统一比较RSA, ElGamal, ECC的性能并绘图。
    """
    print("="*20, "128-bit等效安全级别下的统一性能对比", "="*20)
    print("对比配置: RSA-3072 vs ElGamal-3072 vs ECC-secp256r1")
    print("警告: ElGamal-3072 密钥生成过程将非常缓慢，请耐心等待...")

    # 1. 获取所有算法的性能数据
    rsa_results = run_single_performance_test("RSA", 3072)
    elgamal_results = run_single_performance_test("ElGamal", 3072)
    ecc_results = run_single_performance_test("ECC", "secp256r1")

    # 2. 准备绘图数据
    labels = ['密钥生成', '平均加密', '平均解密']
    rsa_values = [rsa_results['key_gen_time'], rsa_results['encrypt_time'], rsa_results['decrypt_time']]
    elgamal_values = [elgamal_results['key_gen_time'], elgamal_results['encrypt_time'], elgamal_results['decrypt_time']]
    ecc_values = [ecc_results['key_gen_time'], ecc_results['encrypt_time'], ecc_results['decrypt_time']]

    # 3. 开始绘图
    x = np.arange(len(labels))  # 标签位置
    width = 0.25  # 柱子的宽度

    fig, ax = plt.subplots(figsize=(14, 8))
    
    # 绘制三组条形图
    rects1 = ax.bar(x - width, rsa_values, width, label='RSA-3072', color='#ff6347') # 红色系
    rects2 = ax.bar(x, elgamal_values, width, label='ElGamal-3072', color='#90ee90') # 绿色系
    rects3 = ax.bar(x + width, ecc_values, width, label='ECC-secp256r1', color='#87cefa') # 蓝色系

    # 添加标签、标题和图例
    ax.set_ylabel('时间 (毫秒, 对数刻度)')
    ax.set_title('RSA vs ElGamal vs ECC 在128-bit等效安全级别下的性能对比')
    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    ax.set_yscale('log') # 使用对数刻度以更好地显示差异巨大的值
    ax.legend()

    # 在每个柱子顶部添加数值标签
    ax.bar_label(rects1, padding=3, fmt='%.2f')
    ax.bar_label(rects2, padding=3, fmt='%.2f')
    ax.bar_label(rects3, padding=3, fmt='%.2f')

    fig.tight_layout()
    
    # 保存并显示图表
    plt.savefig('all_algorithms_comparison.png')
    plt.show()

# ==============================================================================
# 主程序入口
# ==============================================================================
if __name__ == "__main__":
    # 设置中文字体以正确显示图表
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei'] 
    plt.rcParams['axes.unicode_minus'] = False 
    
    # 运行三者统一对比测试
    compare_all_at_128bit_security()
    
    print("\n\n对比演示已完成！")